{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "248097b8",
   "metadata": {},
   "source": [
    "# EasyEdit Example with **Wise**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "753b8801",
   "metadata": {},
   "source": [
    "In this tutorial, we use `QLoRA` to edit `llama-3.2-3b-instruct` model, we hope this tutorial could help you understand how to use the method QLoRA on LLMs, using the QLoRA method with the llama-3.2-3b-instruct as an example."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b0a1701",
   "metadata": {},
   "source": [
    "## Model Editing\n",
    "\n",
    "Deployed models may still make unpredictable errors. For example, Large Language Models (LLMs) notoriously hallucinate, perpetuate bias, and factually decay, so we should be able to adjust specific behaviors of pre-trained models.\n",
    "\n",
    "**Model editing** aims to adjust an initial base model's $(f_\\theta)$ behavior on the particular edit descriptor $[x_e, y_e]$, such as:\n",
    "- $x_e$: \"Who is the president of the US?\n",
    "- $y_e$: \"Joe Biden.\"\n",
    "\n",
    "efficiently without influencing the model behavior on unrelated samples. The ultimate goal is to create an edited model$(f_\\theta’)$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9717af3a",
   "metadata": {},
   "source": [
    "## 📂 Data Preparation\n",
    "\n",
    "The datasets used can be found in [Google Drive Link](https://drive.google.com/file/d/1YtQvv4WvTa4rJyDYQR2J-uK8rnrt0kTA/view?usp=sharing) (ZsRE)\n",
    "\n",
    "Each dataset contains both an **edit set** and a train set."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cf48075",
   "metadata": {},
   "source": [
    "## Prepare the runtime environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6356ed23",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/mnt/8t/fangjizhan/EasyEdit\n",
      "data\t    examples  multimodal_edit.py   run_wise_editing.sh\n",
      "demo\t    figs      outputs\t\t   tutorial-notebooks\n",
      "Dockerfile  hparams   README.md\t\t   tutorial.pdf\n",
      "easyeditor  LICENSE   requirements.txt\n",
      "edit.py     logs      run_wise_editing.py\n"
     ]
    }
   ],
   "source": [
    "## Clone Repo\n",
    "#!git clone https://github.com/zjunlp/EasyEdit\n",
    "%cd EasyEdit\n",
    "!ls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a104cd71",
   "metadata": {},
   "outputs": [],
   "source": [
    "!apt-get install python3.9\n",
    "!sudo update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 1\n",
    "!sudo update-alternatives --config python3\n",
    "!apt-get install python3-pip\n",
    "%pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b039c94a",
   "metadata": {},
   "source": [
    "## Config Method  Parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d553b513",
   "metadata": {},
   "source": [
    "```python\n",
    "alg_name: \"QLoRA\"\n",
    "model_name: \"./hugging_cache/llama-3.2-3b-instruct\"\n",
    "device: 1\n",
    "\n",
    "# QLoRA specific settings\n",
    "quantization_bit: 4\n",
    "double_quant: true\n",
    "quant_type: \"nf4\" # nf4, fp4， int4, int8\n",
    "\n",
    "# LoRA settings\n",
    "lora_type: \"lora\"  # QLoRA typically uses standard LoRA, not AdaLoRA\n",
    "lora_r: 8\n",
    "lora_alpha: 32\n",
    "lora_dropout: 0.1\n",
    "target_modules: [\"q_proj\", \"v_proj\"]\n",
    "\n",
    "# Training settings\n",
    "num_steps: 50\n",
    "batch_size: 1\n",
    "max_length: 30\n",
    "lr: 5e-3\n",
    "weight_decay: 0.0\n",
    "\n",
    "# Additional settings\n",
    "model_parallel: false\n",
    "\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e9aef0a",
   "metadata": {},
   "source": [
    "## Import models & Run\n",
    "\n",
    "### Edit llama-3.2-3b-instruct on ZsRE with QLoRA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5c0a266f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/mnt/8t/xkw/EasyEdit\n"
     ]
    }
   ],
   "source": [
    "%cd .."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2100450c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from easyeditor import BaseEditor\n",
    "from easyeditor import QLoRAHyperParams\n",
    "\n",
    "prompts = ['Question:What sport does Lionel Messi play? Answer:',\n",
    "                'Question:What role does Cristiano Ronaldo play in football? Answer:',\n",
    "                'Question:Which NBA team does Stephen Curry play for? Answer:']\n",
    "ground_truth = ['football', 'forward', 'Golden State Warriors']\n",
    "target_new = ['basketball', 'defender', 'New York Knicks']\n",
    "subject = ['Lionel Messi', 'Cristiano Ronaldo', 'Stephen Curry']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a0ca0f1d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-12-01 13:25:38,294 - easyeditor.editors.editor - INFO - Instantiating model\n",
      "12/01/2024 13:25:38 - INFO - easyeditor.editors.editor -   Instantiating model\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.004497051239013672,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": null,
       "postfix": null,
       "prefix": "Loading checkpoint shards",
       "rate": null,
       "total": 2,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f0d249ab83cb4c919a4c10c7d767700e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-12-01 13:25:44,163 - easyeditor.editors.editor - INFO - AutoRegressive Model detected, set the padding side of Tokenizer to left...\n",
      "12/01/2024 13:25:44 - INFO - easyeditor.editors.editor -   AutoRegressive Model detected, set the padding side of Tokenizer to left...\n",
      "  0%|          | 0/3 [00:00<?, ?it/s]We detected that you are passing `past_key_values` as a tuple and this is deprecated and will be removed in v4.43. Please use an appropriate `Cache` class (https://huggingface.co/docs/transformers/v4.41.3/en/internal/generation_utils#transformers.Cache)\n",
      "100%|██████████| 3/3 [00:02<00:00,  1.21it/s]\n",
      "  0%|          | 0/3 [00:00<?, ?it/s]`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.\n",
      " 33%|███▎      | 1/3 [00:05<00:11,  5.55s/it]12/01/2024 13:25:52 - INFO - peft.tuners.tuners_utils -   Already found a `peft_config` attribute in the model. This will lead to having multiple adapters in the model. Make sure to know what you are doing!\n",
      " 67%|██████▋   | 2/3 [00:11<00:05,  5.54s/it]12/01/2024 13:25:57 - INFO - peft.tuners.tuners_utils -   Already found a `peft_config` attribute in the model. This will lead to having multiple adapters in the model. Make sure to know what you are doing!\n",
      "100%|██████████| 3/3 [00:16<00:00,  5.54s/it]\n",
      "2024-12-01 13:26:04,025 - easyeditor.editors.editor - INFO - 0 editing: Question:What sport does Lionel Messi play? Answer: -> basketball  \n",
      "\n",
      " {'pre': {'rewrite_acc': [0.0], 'portability': {}}, 'case_id': 0, 'requested_rewrite': {'prompt': 'Question:What sport does Lionel Messi play? Answer:', 'target_new': 'basketball', 'ground_truth': 'football', 'portability': {}, 'locality': {}, 'subject': 'Lionel Messi'}, 'post': {'rewrite_acc': [0.0], 'locality': {}, 'portability': {}}}\n",
      "12/01/2024 13:26:04 - INFO - easyeditor.editors.editor -   0 editing: Question:What sport does Lionel Messi play? Answer: -> basketball  \n",
      "\n",
      " {'pre': {'rewrite_acc': [0.0], 'portability': {}}, 'case_id': 0, 'requested_rewrite': {'prompt': 'Question:What sport does Lionel Messi play? Answer:', 'target_new': 'basketball', 'ground_truth': 'football', 'portability': {}, 'locality': {}, 'subject': 'Lionel Messi'}, 'post': {'rewrite_acc': [0.0], 'locality': {}, 'portability': {}}}\n",
      "2024-12-01 13:26:04,786 - easyeditor.editors.editor - INFO - 1 editing: Question:What role does Cristiano Ronaldo play in football? Answer: -> defender  \n",
      "\n",
      " {'pre': {'rewrite_acc': [0.0], 'portability': {}}, 'case_id': 1, 'requested_rewrite': {'prompt': 'Question:What role does Cristiano Ronaldo play in football? Answer:', 'target_new': 'defender', 'ground_truth': 'forward', 'portability': {}, 'locality': {}, 'subject': 'Cristiano Ronaldo'}, 'post': {'rewrite_acc': [0.0], 'locality': {}, 'portability': {}}}\n",
      "12/01/2024 13:26:04 - INFO - easyeditor.editors.editor -   1 editing: Question:What role does Cristiano Ronaldo play in football? Answer: -> defender  \n",
      "\n",
      " {'pre': {'rewrite_acc': [0.0], 'portability': {}}, 'case_id': 1, 'requested_rewrite': {'prompt': 'Question:What role does Cristiano Ronaldo play in football? Answer:', 'target_new': 'defender', 'ground_truth': 'forward', 'portability': {}, 'locality': {}, 'subject': 'Cristiano Ronaldo'}, 'post': {'rewrite_acc': [0.0], 'locality': {}, 'portability': {}}}\n",
      "2024-12-01 13:26:05,548 - easyeditor.editors.editor - INFO - 2 editing: Question:Which NBA team does Stephen Curry play for? Answer: -> New York Knicks  \n",
      "\n",
      " {'pre': {'rewrite_acc': [0.6666666666666666], 'portability': {}}, 'case_id': 2, 'requested_rewrite': {'prompt': 'Question:Which NBA team does Stephen Curry play for? Answer:', 'target_new': 'New York Knicks', 'ground_truth': 'Golden State Warriors', 'portability': {}, 'locality': {}, 'subject': 'Stephen Curry'}, 'post': {'rewrite_acc': [0.0], 'locality': {}, 'portability': {}}}\n",
      "12/01/2024 13:26:05 - INFO - easyeditor.editors.editor -   2 editing: Question:Which NBA team does Stephen Curry play for? Answer: -> New York Knicks  \n",
      "\n",
      " {'pre': {'rewrite_acc': [0.6666666666666666], 'portability': {}}, 'case_id': 2, 'requested_rewrite': {'prompt': 'Question:Which NBA team does Stephen Curry play for? Answer:', 'target_new': 'New York Knicks', 'ground_truth': 'Golden State Warriors', 'portability': {}, 'locality': {}, 'subject': 'Stephen Curry'}, 'post': {'rewrite_acc': [0.0], 'locality': {}, 'portability': {}}}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Metrics Summary:  {'pre': {'rewrite_acc': 0.2222222222222222}, 'post': {'rewrite_acc': 0.0}}\n"
     ]
    }
   ],
   "source": [
    "\n",
    "hparams = QLoRAHyperParams.from_hparams('./hparams/QLoRA/llama3.2-3b.yaml')\n",
    "\n",
    "editor = BaseEditor.from_hparams(hparams)\n",
    "metrics, edited_model, _ = editor.edit(\n",
    "    prompts=prompts,\n",
    "    ground_truth=ground_truth,\n",
    "    target_new=target_new,\n",
    "    subject=subject,\n",
    "    sequential_edit=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe56c3a1",
   "metadata": {},
   "source": [
    "* edit_data: editing instance in edit set.\n",
    "* loc_data: used to provide xi in Equation 5, sampled from the train set.\n",
    "* sequential_edit: whether to enable sequential editing (should be set to True except when T=1).\n",
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3db4502",
   "metadata": {},
   "source": [
    "### Reliability Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cc703696",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.00561976432800293,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": null,
       "postfix": null,
       "prefix": "Loading checkpoint shards",
       "rate": null,
       "total": 2,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "159a756fceab47219bf1c783fc4cd80a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "from transformers import LlamaForCausalLM\n",
    "\n",
    "device = 0\n",
    "model = LlamaForCausalLM.from_pretrained('./hugging_cache/llama-3.2-3b-instruct').to(f'cuda:{device}')  \n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained('./hugging_cache/llama-3.2-3b-instruct',trust_remote_code=True)\n",
    "tokenizer.pad_token_id = tokenizer.eos_token_id\n",
    "tokenizer.padding_side='left'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c2acf594",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prompt: Question:What sport does Lionel Messi play? Answer:\n",
      "Pre-Edit  Output: Football (Soccer)\n",
      "What sport does Lionel Messi\n",
      "Post-Edit Output:  What is this? The'mea' that\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Prompt: Question:What role does Cristiano Ronaldo play in football? Answer:\n",
      "Pre-Edit  Output:  Cristiano Ronaldo is a professional footballer who plays as\n",
      "Post-Edit Output:  Cristiano Ronaldo plays the position of the 'Number \n",
      "----------------------------------------------------------------------------------------------------\n",
      "Prompt: Question:Which NBA team does Stephen Curry play for? Answer:\n",
      "Pre-Edit  Output:  The Golden State Warriors.\n",
      "The NBA is the National\n",
      "Post-Edit Output:  His last 5, 2nd, \n",
      "----------------------------------------------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "correct_prompts = ['Question:What sport does Lionel Messi play? Answer:',\n",
    "                'Question:What role does Cristiano Ronaldo play in football? Answer:',\n",
    "                'Question:Which NBA team does Stephen Curry play for? Answer:']\n",
    "\n",
    "# ground_truth = ['football', 'forward', 'Golden State Warriors']\n",
    "# target_new = ['basketball', 'defender', 'New York Knicks']\n",
    "batch = tokenizer(correct_prompts, return_tensors='pt', padding=True)\n",
    "\n",
    "pre_edit_outputs = model.generate(\n",
    "    input_ids=batch['input_ids'].to(model.device),\n",
    "    attention_mask=batch['attention_mask'].to(model.device),\n",
    "    pad_token_id = tokenizer.eos_token_id,\n",
    "    max_new_tokens=10\n",
    ")\n",
    "post_edit_outputs = edited_model.generate(\n",
    "    input_ids=batch['input_ids'].to(edited_model.device),\n",
    "    attention_mask=batch['attention_mask'].to(edited_model.device),\n",
    "    pad_token_id = tokenizer.eos_token_id,\n",
    "    max_new_tokens=10\n",
    ")\n",
    "\n",
    "max_length = batch['input_ids'].shape[-1]\n",
    "for i in range(len(correct_prompts)):\n",
    "    print(f'Prompt: {correct_prompts[i]}')\n",
    "    print(f'Pre-Edit  Output: {tokenizer.decode( pre_edit_outputs[i][max_length:], skip_special_tokens=True)}')\n",
    "    print(f'Post-Edit Output: {tokenizer.decode(post_edit_outputs[i][max_length:], skip_special_tokens=True)}')\n",
    "    print('--'*50 )\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43528147",
   "metadata": {},
   "source": [
    "### Generalization test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4074b583",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages/transformers/tokenization_utils_base.py:2888: UserWarning: `max_length` is ignored when `padding`=`True` and there is no truncation strategy. To pad to max length, use `padding='max_length'`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prompt: Question:What sports is Messi good at? Answer:\n",
      "Pre-Edit  Output:  Messi is good at soccer, or football as it\n",
      "Post-Edit Output:  Who would have thought it was the 'Who'?\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Prompt: Question:What position does Cristiano Ronaldo hold in the sport of football? Answer:\n",
      "Pre-Edit  Output: Striker.\n",
      "What position does Cristiano Ronaldo hold in\n",
      "Post-Edit Output:  He is an attacking midfielder, but he has been\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Prompt: Question:Which city does Stephen Curry currently working in? Answer:\n",
      "Pre-Edit  Output: He currently works in Oakland, California. \n",
      "However\n",
      "Post-Edit Output:  It is actually in New York City?\n",
      "----------------------------------------------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "generation_prompts =['Question:What sports is Messi good at? Answer:',\n",
    "                'Question:What position does Cristiano Ronaldo hold in the sport of football? Answer:',\n",
    "                'Question:Which city does Stephen Curry currently working in? Answer:']\n",
    "\n",
    "batch = tokenizer(generation_prompts , return_tensors='pt', padding=True, max_length=30)\n",
    "\n",
    "pre_edit_outputs = model.generate(\n",
    "    input_ids=batch['input_ids'].to(model.device),\n",
    "    attention_mask=batch['attention_mask'].to(model.device),\n",
    "    pad_token_id = tokenizer.eos_token_id,\n",
    "    max_new_tokens=10\n",
    ")\n",
    "post_edit_outputs = edited_model.generate(\n",
    "    input_ids=batch['input_ids'].to(edited_model.device),\n",
    "    attention_mask=batch['attention_mask'].to(edited_model.device),\n",
    "    pad_token_id = tokenizer.eos_token_id,\n",
    "    max_new_tokens=10\n",
    ")\n",
    "\n",
    "max_length = batch['input_ids'].shape[-1]\n",
    "for i in range(len(generation_prompts)):\n",
    "    print(f'Prompt: {generation_prompts[i]}')\n",
    "    print(f'Pre-Edit  Output: {tokenizer.decode( pre_edit_outputs[i][max_length:], skip_special_tokens=True)}')\n",
    "    print(f'Post-Edit Output: {tokenizer.decode(post_edit_outputs[i][max_length:], skip_special_tokens=True)}')\n",
    "    print('--'*50 )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d4c3779",
   "metadata": {},
   "source": [
    "### Locality test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f21404e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prompt: Question:What sport does Kylian Mbappé play? Answer:\n",
      "Pre-Edit  Output: Football\n",
      "What sport does Kylian Mbappé\n",
      "Post-Edit Output:  He is known to play for, and be a\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Prompt: Question:What role does Thierry Henry play in football? Answer:\n",
      "Pre-Edit  Output:  Thierry Henry is a former professional footballer who\n",
      "Post-Edit Output:  He is a great striker... or what?\n",
      "To\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Prompt: Question:Which NBA team does Jordan play for? Answer:\n",
      "Pre-Edit  Output:  The Chicago Bulls.\n",
      "Answer: The Chicago Bulls.\n",
      "\n",
      "Post-Edit Output:  The 'S' in 'S'\n",
      "----------------------------------------------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "locality_prompts = ['Question:What sport does Kylian Mbappé play? Answer:',\n",
    "                'Question:What role does Thierry Henry play in football? Answer:',\n",
    "                'Question:Which NBA team does Jordan play for? Answer:']\n",
    "\n",
    "batch = tokenizer(locality_prompts, return_tensors='pt', padding=True)\n",
    "\n",
    "pre_edit_outputs = model.generate(\n",
    "    input_ids=batch['input_ids'].to(model.device),\n",
    "    attention_mask=batch['attention_mask'].to(model.device),\n",
    "    pad_token_id = tokenizer.eos_token_id,\n",
    "    max_new_tokens=10\n",
    ")\n",
    "post_edit_outputs = edited_model.generate(\n",
    "    input_ids=batch['input_ids'].to(edited_model.device),\n",
    "    attention_mask=batch['attention_mask'].to(edited_model.device),\n",
    "    pad_token_id = tokenizer.eos_token_id,\n",
    "    max_new_tokens=10\n",
    ")\n",
    "\n",
    "max_length = batch['input_ids'].shape[-1]\n",
    "for i in range(len(locality_prompts )):\n",
    "    print(f'Prompt: {locality_prompts[i]}')\n",
    "    print(f'Pre-Edit  Output: {tokenizer.decode( pre_edit_outputs[i][max_length:], skip_special_tokens=True)}')\n",
    "    print(f'Post-Edit Output: {tokenizer.decode(post_edit_outputs[i][max_length:], skip_special_tokens=True)}')\n",
    "    print('--'*50 )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "EasyEdit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
